//! 在Response中插入origin

use axum::{
    http::{header, Request, StatusCode},
    response::{IntoResponse, Response},
    BoxError,
};
use tower::{Layer, Service};

use crate::middleware::{pin_box_future, MiddleError, PinBoxFuture};

#[derive(Clone)]
pub struct ResponseWithOriginLayer;

impl ResponseWithOriginLayer {
    pub fn new() -> Self {
        ResponseWithOriginLayer
    }
}

impl<S> Layer<S> for ResponseWithOriginLayer {
    type Service = ResponseWithOrigin<S>;

    fn layer(&self, inner: S) -> Self::Service {
        ResponseWithOrigin { inner }
    }
}

#[derive(Debug, Clone)]
pub struct ResponseWithOrigin<S> {
    inner: S,
}

impl<B, S> Service<Request<B>> for ResponseWithOrigin<S>
where
    S: Service<Request<B>>,
    S::Error: Into<BoxError>,
    S::Future: Send + 'static,
    S::Response: IntoResponse,
    B: Send + 'static,
{
    type Response = Response;
    type Error = BoxError;
    type Future = PinBoxFuture<Result<Self::Response, Self::Error>>;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx).map_err(Into::into)
    }

    //noinspection RsUnresolvedReference
    //noinspection ALL
    fn call(&mut self, req: Request<B>) -> Self::Future {
        let origin = req.headers().get(header::ORIGIN).cloned();
        let req = self.inner.call(req);
        pin_box_future(async move {
            let origin = origin.ok_or::<Self::Error>(
                MiddleError(StatusCode::BAD_REQUEST, "header中丢失origin".to_string()).into(),
            )?;

            let mut res = req
                .await
                .map_err(Into::into)
                .map(IntoResponse::into_response)?;
            // 每个响应都要带上ACCESS_CONTROL_ALLOW_ORIGIN
            let _ = res
                .headers_mut()
                .insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, origin.clone());
            Ok(res)
        })
    }
}
